﻿using System;
using System.Collections.Generic;
//using System.Linq;
//using System.Text;
using System.Diagnostics;
using Allegro.Framework;
using Allegro.Integrator;
using Allegro.Integrator.RungeKutta;

//using Allegro.Integrator.RungeKutta;

namespace Allegro.Mathlib
{
    public enum IntegrationMethod : int
    {
        AdaptiveRungeKutta = 0, // default
        BulirschStoer = 1
    }
    /// <summary>
    /// Delegate for pushing results from the solver to the caller
    /// </summary>
    /// <param name="x"></param>
    /// <param name="y"></param>
    public delegate void ProcessResult(double x, double[] y);

    /// <summary>
    /// Interception of the integration at select times during integration, 
    /// e.g. prior to taking a step (pre-step interception), after taking a step (post-step interception).
    /// Delegate should return false if integration should be terminated, false if it should be continued.
    /// </summary>
    /// <param name="x">Current x</param>
    /// <param name="y">Current y vector</param>
    /// <returns></returns>
    public delegate bool InterceptIntegration(double x, double[] y);

    public class IntegrationSettings
    {
        double _hTrial;
        bool _hTrialSet;
        double _eps;
        //bool _accuracySet;
        double _hMin;
        //bool _hMinSet;
        double _x1;
        double _x2;
        int _nMaxSteps;
        private IntegrationMethod _method;
        
        internal IntegrationSettings(int maxNumberOfSteps, double accuracy)
        {
            _eps = accuracy;
            _nMaxSteps = maxNumberOfSteps;
            _hTrial = 0;
            _hMin = 0;
            nok = nbad = 0;
            _hTrialSet = false;
            //_hMinSet = false;
            _x2 = _x1 = 0;
            _method = IntegrationMethod.AdaptiveRungeKutta;
        }
        public double x2
        {
            get { return _x2; }
            set { _x2 = value; }
        }
        public double x1
        {
            get { return _x1; }
            set { _x1 = value; }
        }
        /// <summary>
        /// Max. number of steps
        /// </summary>
        public int nMaxStep
        {
            get { return _nMaxSteps; }
            set { _nMaxSteps = value; }
        }
        /// <summary>
        /// Trial step size. Default is (x2-x1)/nMaxStep
        /// </summary>
        public double hTrial
        {
            get 
            {
                if (!_hTrialSet)
                    return (_x2 - _x1) / _nMaxSteps;
                else
                    return _hTrial; 
            }
            set 
            {
                _hTrial = value;
                _hTrialSet = true;
            }
        }
        /// <summary>
        /// Min. step size
        /// </summary>
        public double hMin
        {
            get { return _hMin; }
            set 
            {
                _hMin = value;
                //_hMinSet = true;
            }
        }

        /// <summary>
        /// Relative accuracy
        /// Default value = 10^-8
        /// </summary>
        public double eps
        {
            get { return _eps; }
            set { _eps = value; }
        }
        /// <summary>
        /// Number of ok steps. Set when computation is completed. 
        /// </summary>
        public int nok;
        /// <summary>
        /// Number of bad steps taken. Set when computation is completed
        /// </summary>
        public int nbad;

        /// <summary>
        /// Integration method. Default is AdaptiveRungeKutta
        /// </summary>
        public IntegrationMethod Method
        {
            get{ return _method;}
            set{ _method = value;}
        }
    }

    /// <summary>
    /// Driver for ODE integration
    /// </summary>
    public class BasicODESolver : IODESolver
    {
        IntegrationSettings _options;
        /// <summary>
        /// Default value for maximum number of steps to be taken
        /// </summary>
        public const int MAXSTP = 10000;
        public const double DefaultAccuracy = 1.0e-8;
        //public const double TINY = 1.0e-30;

        //double _dxsav; // Result save interval
        int _nMaxStp;

        /// <summary>
        /// Initialize the solver. Define the maximum number of integration steps to be taken as well 
        /// as the relative accuracy. Default accuracy is 10^-7.
        /// </summary>
        /// <param name="integrator"></param>
        public BasicODESolver(int nMaxStp, double accuracy)
        {
            _options = new IntegrationSettings(nMaxStp, accuracy);
            //_dxsav = 0;
            _nMaxStp = nMaxStp;
        }
        /// <summary>
        /// Initialize the solver. Define the maximum number of integration steps to be taken.
        /// </summary>
        /// <param name="nMaxStp"></param>
        public BasicODESolver(int nMaxStp)
            : this(nMaxStp, DefaultAccuracy)
        {
        }
        /// <summary>
        /// Initialize the solver. The maximum number of integration steps to be taken is the
        /// default value, which is 10000.
        /// </summary>
        public BasicODESolver() : this(MAXSTP)
        {
        }

        /// <summary>
        /// Result save interval
        /// </summary>
        //public double ResultSaveInterval
        //{
        //    get { return _dxsav; }
        //    set { _dxsav = value; }
        //}

        /// <summary>
        /// Runge-Kutta driver with adaptive stepsize control. Integrate starting values ystart[1..n]
        /// from x1 with accuracy eps. The condition for stopping the integration is delegated to the exitIf delegate, which can be supplied by the caller.
        /// This enables the caller to control when the computation should end. If exitIf==null, the computation will end when x >= x2 or when the
        /// number of integration steps exceeds the maximum number of steps set.
        /// h1 should be set as a guessed first stepsize, hmin as the minimum allowed stepsize
        /// (can be zero). On output nok and nbad are the number of good and bad (but retried and fixed) steps taken, and
        /// ystart is replaced by values at the end of the integration interval. derivs is the user-supplied method for calculating
        /// the right-hand side derivative, while integrator is the integrator to be used to calculate each step. resultHandler is the
        /// user-supplied result set delegate for handling the results from each integration step. resultHandler will be invoked at each step.
        /// </summary>
        /// <param name="integrator"></param>
        /// <param name="exitIf">Exit condition</param>
        /// <param name="resultHandler"></param>
        /// <param name="derivs"></param>
        /// <param name="ystart"></param>
        /// <param name="x1"></param>
        /// <param name="x2">Max. value of x. Is also used to set the direction of integration. If (x2-x1)&lt; 0, the step size will be negative.</param>
        /// <param name="eps"></param>
        /// <param name="h1">Trial step size</param>
        /// <param name="hmin">Min. step size</param>
        /// <param name="nok">Number of good steps taken</param>
        /// <param name="nbad">Number of bad (but retried and fixed) steps taken</param>
        public void Solve(IODEIntegrator integrator, InterceptIntegration preStepIntercept, InterceptIntegration postStepIntercept, 
            ProcessResult resultHandler, Derivative derivs, ComputeValues values, double[] ystart, double x1, double x2,
            double eps, double h1, double hmin, out int nok, out int nbad)
        {
            double x, hnext, hdid, h;
            //double[] yscal, y, dydx;
            double[] y;
            int nvar = ystart.GetLength(0);

            // Validate input
            if (nvar <= 0)
                ErrorHandler.Error("BasicODESolver.Solve: Illegal nvar value, must be > 0");
            if (resultHandler == null)
                ErrorHandler.Error("BasicODESolver.Solve: No resultHandler");
            if (integrator == null)
                ErrorHandler.Error("BasicODESolver.Solve: No integrator");

            integrator.Accuracy = eps;

            nok = 0;
            nbad = 0;

            //yscal = new double[nvar];
            y = new double[nvar];
            //dydx = new double[nvar];
            x = x1;
            h = (x2 - x1 > 0 ? Math.Abs(h1) : -Math.Abs(h1));

            for (int i = 0; i < nvar; i++)
            {
                y[i] = ystart[i];
            }
            int nTotalSteps = 0;
            // Take at most _nMaxStp steps
            for (int nstp = 0; nstp < _nMaxStp; nstp++)
            {
                if (preStepIntercept != null && !preStepIntercept(x, y))
                    break;

                // Store intermediate results
                resultHandler(x, y);

                // If stepsize can overshoot, decrease
                if ((x + h - x2) * (x + h - x1) > 0)
                    h = x2 - x;
                // Make trial integration step
                integrator.Integrate(ref x, ref y, h, out hdid, out hnext, derivs, values);
                if (hdid == h)
                    nok++;
                else
                    nbad++;
                // Are we done?
                if (postStepIntercept != null && !postStepIntercept(x, y))
                {
                    resultHandler(x, y);
                    nTotalSteps = nstp;
                    break;
                }
                if ((x-x2)*(x2-x1) >= 0)
                {
                    resultHandler(x, y);
                    nTotalSteps = nstp;
                    break;
                }
                if (Math.Abs(hnext) <= hmin)
                    ErrorHandler.Error(string.Format("Step size {0} too small at x={1}", hnext, x));
                h = hnext;
                if(nstp == _nMaxStp-1)
                    ErrorHandler.Error(string.Format("Too many integration steps taken. End at x = {0}", x));
            }
            Trace.WriteLine(string.Format("Number of steps taken: {0}", nTotalSteps));
        }

        #region IODESolver Members

        public void Solve(IStepComputation stepComputationProvider, IResultset resultset, double x1, double x2, double[] ystart)
        {
            int nDim = ystart.GetLength(0);
            
            IODEIntegrator integrator;
            if (Settings.Method == IntegrationMethod.BulirschStoer)
            {
                integrator = new BulirschStoerIntegrator(nDim);
            }
            else
            {
                ButcherTableau rkTableau = new RKTableau_56_Verner02();
                integrator = new GenericRungeKuttaIntegrator(nDim, rkTableau);
                //integrator = new RungeKuttaIntegrator(nDim);
            }

            Settings.x1 = x1;
            Settings.x2 = x2;

            Solve(integrator, stepComputationProvider.PrestepComputation, stepComputationProvider.PoststepComputation,
                resultset.AddData, stepComputationProvider.DyDx, stepComputationProvider.ComputeValues, 
                ystart, x1, x2, _options.eps, _options.hTrial,
                _options.hMin, out _options.nok, out _options.nbad);
        }

        public IntegrationSettings Settings 
        {
            get { return _options; }
            set { _options = value; }
        }

        #endregion
    }
}
